{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from langchain.chains import ConversationChain\n",
    "from langchain.chains.llm import LLMChain\n",
    "from langchain.chains.router import MultiPromptChain\n",
    "from langchain_community.llms import Tongyi\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.chains.router.llm_router import LLMRouterChain, RouterOutputParser\n",
    "from langchain.chains.router.multi_prompt_prompt import MULTI_PROMPT_ROUTER_TEMPLATE\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "physics_template = '''\n",
    "你是一位非常聪明的物理学教授。你擅长以简明易懂的方式回答物理方面的问题。当你不知道一个问题的答案时，你会诚实地承认。\n",
    "\n",
    "问题：{input}\n",
    "'''\n",
    "\n",
    "math_template = '''\n",
    "你是一位出色的数学家，擅长解决各种数学问题。你能够将复杂的问题分解成易于处理的部分，并逐一解答，然后综合这些部分提供完整的答案。\n",
    "\n",
    "问题：{input}\n",
    "'''\n",
    "\n",
    "history_template = '''\n",
    "你是一位资深的历史学家，对不同历史时期的人物、事件和背景有深入的了解。你具备分析、反思、辩论和评估历史事件的能力，尊重历史证据并能利用它们支持你的解释和判断。\n",
    "\n",
    "问题：{input}\n",
    "'''\n",
    "\n",
    "computerscience_template = '''\n",
    "你是一位经验丰富的计算机科学家，精通各种编程语言和计算机技术。你能够清晰地解释复杂的计算机科学概念，并提供实用的解决方案。\n",
    "\n",
    "问题：{input}\n",
    "'''\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_infos = [\n",
    "    {\n",
    "        \"name\": \"physics\",\n",
    "        \"description\": \"适合回答关于物理的问题\",\n",
    "        \"prompt_template\": physics_template\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"math\",\n",
    "        \"description\": \"适合回答关于数学的问题\",\n",
    "        \"prompt_template\": math_template\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"history\",\n",
    "        \"description\": \"适合回答关于历史的问题\",\n",
    "        \"prompt_template\": history_template\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"computer_science\",\n",
    "        \"description\": \"适合回答关于计算机科学的问题\",\n",
    "        \"prompt_template\": computerscience_template\n",
    "    }\n",
    "]\n",
    "\n",
    "# 初始化与 OpenAI 大型语言模型的连接\n",
    "llm = Tongyi(dashscope_api_key=os.getenv(\"DASHSCOPE_API_KEY\"), model_name=\"qwen-plus\")\n",
    "\n",
    "# 创建 destination_chains 字典，存储各学科对应的子链\n",
    "destination_chains = {}\n",
    "for p_info in prompt_infos:\n",
    "    name = p_info[\"name\"]\n",
    "    prompt_template = p_info[\"prompt_template\"]\n",
    "    \n",
    "    # 使用 PromptTemplate 构造函数创建 PromptTemplate 实例\n",
    "    prompt = PromptTemplate(\n",
    "        template=prompt_template,\n",
    "        input_variables=[\"input\"]\n",
    "    )\n",
    "    \n",
    "    # 创建 LLMChain 实例\n",
    "    chain = LLMChain(llm=llm, prompt=prompt)\n",
    "    destination_chains[name] = chain\n",
    "\n",
    "# 打印各目标链的信息\n",
    "destinations = [f\"{p['name']}: {p['description']}\" for p in prompt_infos]\n",
    "destinations_str = \"\\n\".join(destinations)\n",
    "print(destinations_str)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains.router.llm_router import LLMRouterChain, RouterOutputParser\n",
    "from langchain.chains.router.multi_prompt_prompt import MULTI_PROMPT_ROUTER_TEMPLATE\n",
    "# 创建默认链\n",
    "default_chain = ConversationChain(llm=llm, output_key=\"text\")\n",
    "\n",
    "# 配置路由器模板\n",
    "router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(\n",
    "    destinations=destinations_str\n",
    ")\n",
    "\n",
    "# 创建路由器提示\n",
    "router_prompt = PromptTemplate(\n",
    "    template=router_template,\n",
    "    input_variables=[\"input\"],\n",
    "    output_parser=RouterOutputParser()\n",
    ")\n",
    "\n",
    "# 创建路由器链\n",
    "router_chain = LLMRouterChain.from_llm(llm, router_prompt)\n",
    "\n",
    "# 创建多提示链\n",
    "chain = MultiPromptChain(\n",
    "    router_chain=router_chain,\n",
    "    destination_chains=destination_chains,\n",
    "    default_chain=default_chain,\n",
    "    verbose=True\n",
    ")\n",
    "\n",
    "# 运行示例\n",
    "print(chain.run(\"什么是牛顿定律？\"))  # 应路由到物理链\n",
    "print(chain.run(\"三国争霸中，曹操为什么能够胜出？\"))  # 应路由到历史链\n",
    "print(chain.run(\"感冒发烧吃什么药合适？\"))  # 应路由到默认链"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chatchat",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
